import torch
import torch.nn as nn
import torch.nn.functional as F
from einops.layers.torch import Rearrange

class CNNFeatureExtractor(nn.Module):
    def __init__(self, emb_size=20):
        super().__init__()
        self.shallownet = nn.Sequential(
            nn.Conv2d(1, 20, (1, 25), (1, 1)),
            nn.ReLU(),
            nn.Conv2d(20, 20, (22, 1), (1, 1)),
            nn.BatchNorm2d(20),
            nn.ReLU(),
            nn.AvgPool2d((1, 45), (1, 15)),
            nn.Dropout(0.5),
            nn.Conv2d(20, emb_size, (1, 1), stride=(1, 1)),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),  # 全局平均池化
            Rearrange('b e (h) (w) -> b (e h w)')
        )

    def forward(self, x):
        return self.shallownet(x)

class CNNSimple(nn.Module):
    def __init__(self, emb_size=20, n_classes=2, device=None):
        super().__init__()
        self.feature_extractor = CNNFeatureExtractor(emb_size=emb_size)
        self.classifier = nn.Linear(emb_size, n_classes)
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.sigmoid = nn.Sigmoid()  # 使用sigmoid来确保输出范围在[0, 1]

    def forward(self, x, return_features=False):
        features = self.feature_extractor(x)
        
        logits = self.classifier(features)
        probs = self.sigmoid(logits)  # 使用sigmoid获取[0,1]之间的输出
        if return_features:
            return probs  # 如果需要特征，直接返回
        return torch.round(probs)  # 应用0.5阈值进行四舍五入，确保输出为0或1

